package OceanModel;

import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

public interface Dynamics {
	// Interface for a class that calculates model derivatives
	public void calculateDerivatives(HashMap<String, GriddedVariable> variables);
	public HashMap<String, double[][][]> getDerivatives();
}

class simpleDynamics implements Dynamics {
	// Class that implements a simplified subset of the Navier-Stokes fluid equations.
	
	Parameters params;
	HashMap<String, double[][][]> derivatives;
	HashMap<String, GriddedVariable> variables;
	  
	// Objects that calculate model derivatives due to various processes.
	Continuity continuity;
	Coriolis coriolis;
	Pressure pressure;
	Viscosity viscosity;
	Forcing forcing;

	//------------------------
	  
	simpleDynamics(Parameters inputParams, 
			       HashMap<String, GriddedVariable> inputVariables, 
			       String[] derivativeNameArray) {
		params = inputParams;
		variables = inputVariables;
	    
		derivatives = new HashMap<String, double[][][]>();

		// Set up derivative variables using the names in the derivativeNameArray
		for (int i=0; i < derivativeNameArray.length; i++) {
			GriddedVariable var = variables.get( derivativeNameArray[i] );
	      
			double[][][] newDerivative = new double[var.nz][var.ny][var.nx];

			derivatives.put( derivativeNameArray[i], newDerivative);
		}
	    
		// Create physics processes objects
		continuity = new Continuity(variables, derivatives, params);
		coriolis = new Coriolis(variables, derivatives, params);
		forcing = new Forcing(variables, derivatives, params);
		pressure = new Pressure(variables, derivatives, params);
		viscosity = new Viscosity(variables, derivatives, params);
	}

	//------------------------
	
	public void calculateDerivatives(HashMap<String, GriddedVariable> variables) {
		// zero out all the derivatives variables
		zero_arrays(variables, derivatives, params);

		// Calculate model variable derivatives by calling each process class.
		continuity.calculate(variables, derivatives, params);
		coriolis.calculate(variables, derivatives, params);
//		forcing.calculate(variables, derivatives, params);		
		pressure.calculate(variables, derivatives, params);
		viscosity.calculate(variables, derivatives, params);

		// Zero out derivatives on land points
		mask_variables(variables, derivatives, params);

	}

	//------------------------
	  
	public HashMap<String, double[][][]> getDerivatives() {
		// derivatives getter
		return derivatives;
	}  

	//------------------------

	void zero_arrays(HashMap<String, GriddedVariable> variables, 
			         HashMap<String, double[][][]> derivatives, 
			         Parameters params) {
		// zero out all the derivatives variables

	    for (Map.Entry<String, double[][][]> deriv : derivatives.entrySet()) {
			GriddedVariable currentVar = variables.get( deriv.getKey() );
			double[][][] deriv_array = deriv.getValue();

			for (int k=0; k < currentVar.nz; k++) {
				for (int j=0; j < currentVar.ny; j++) {
					for (int i=0; i < currentVar.nx; i++) {
						deriv_array[k][j][i] = 0.;
					}
				}
			}	            
		} // end for

	}
	  
	//------------------------
	
	void mask_variables(HashMap<String, GriddedVariable> variables, 
			            HashMap<String, double[][][]> derivatives,
			            Parameters params) {

		// zero out all the derivatives variables in the model grid halo

	    for (Map.Entry<String, double[][][]> deriv : derivatives.entrySet()) {
			GriddedVariable currentVar = variables.get( deriv.getKey() );

			for (int k=0; k < currentVar.nz; k++) {
				for (int j=0; j < currentVar.ny; j++) {
					currentVar.data[k][j][0] = 0.;
					currentVar.data[k][j][currentVar.nx-1] = 0.;
				}
				for (int i=0; i < currentVar.nx; i++) {
					currentVar.data[k][0][i] = 0.;
					currentVar.data[k][currentVar.ny-1][i] = 0.;
				}
			}

			if (currentVar.gridType == 'u') {
				for (int k=0; k < currentVar.nz; k++) {
					for (int j=0; j < currentVar.ny; j++) {
						currentVar.data[k][j][1] = 0.;
						currentVar.data[k][j][currentVar.nx-2] = 0.;
					}
				}
			}
			if (currentVar.gridType == 'v') {
				for (int k=0; k < currentVar.nz; k++) {
					for (int i=0; i < currentVar.nx; i++) {
						currentVar.data[k][1][i] = 0.;
						currentVar.data[k][currentVar.ny-2][i] = 0.;
					}
				}
			}
		} // end while
		
	}
}

	//=================================================

class Continuity {
	
	// Calculates layer thickness changes due to velocity divergence
	  
	HashMap<String, GriddedVariable> variables;
	HashMap<String, double[][][]> derivatives;
	Parameters params;
	  
	Continuity(HashMap<String, GriddedVariable> variables, 
			   HashMap<String, double[][][]> derivatives, 
			   Parameters params) {
		this.variables = variables;
		this.derivatives = derivatives;
		this.params = params;
	}
	  
	//-------------------------------------------
	  
	public void calculate(HashMap<String, GriddedVariable> variables, 
			              HashMap<String, double[][][]> derivatives, 
			              Parameters params) {

		double[][][] data_du_dx, data_dv_dy;
		double[][][] deta_dt = (double[][][]) derivatives.get("eta");
		GriddedVariable h = variables.get("h");
		GriddedVariable u = variables.get("u");
		GriddedVariable v = variables.get("v");
	        
		data_du_dx = u.d_dx().data;
		data_dv_dy = v.d_dy().data;
		for (int k=0; k < h.nz; k++) {
			for (int j=0; j < h.ny; j++) {
				for (int i=0; i < h.nx; i++) {
					deta_dt[k][j][i] = -(data_du_dx[k][j][i] + data_dv_dy[k][j][i])*h.data[k][j][i];
				}
			}
		}
	}
	  
}

//=================================================

class Coriolis {
	double f;
	double beta;  

	// Calculates the coriolis force

	
	Coriolis(HashMap<String, GriddedVariable> variables, 
			 HashMap<String, double[][][]> derivatives, 
			 Parameters params) {
		f = params.f;
		beta = params.beta;
	}

	//-----------------------------------------------------------------------
	  
	public void calculate(HashMap<String, GriddedVariable> variables, 
			              HashMap<String, double[][][]> derivatives, 
			              Parameters params) {
		double y;
		double[][][] dataU, dataV;

		double[][][] du_dt = (double[][][]) derivatives.get("u");
		double[][][] dv_dt = (double[][][]) derivatives.get("v");

		GriddedVariable u = variables.get("u");
		GriddedVariable v = variables.get("v");
	    
		dataV = v.onGrid('u').data;
		for (int k=0; k < u.nz; k++) {
			for (int j=0; j < u.ny; j++) {
				for (int i=0; i < u.nx; i++) {      
					y = ((double)j)*u.dy;
					du_dt[k][j][i] += (f + beta*y) * dataV[k][j][i];
				}
			}
		}
	        
		dataU = u.onGrid('v').data;
		for (int k=0; k < v.nz; k++) {
			for (int j=0; j < v.ny; j++) {
				for (int i=0; i < v.nx; i++) {      
					y = ((double)j - .5)*v.dy;
					dv_dt[k][j][i] -= (f + beta*y) * dataU[k][j][i];
				}
			}
		}
	    
	}
	  
}

//=================================================

class Forcing {

	// Calculates wind forcing

	GriddedVariable taux;
	GriddedVariable tauy;

	Forcing(HashMap<String, GriddedVariable> variables, 
			HashMap<String, double[][][]> derivatives, 
			Parameters params) {
		taux = new GriddedVariable(params, 'u');
		tauy = new GriddedVariable(params, 'v');
		double y;
		
		for (int j=0; j < taux.ny; j++) {
			for (int i=0; i < taux.nx; i++) {      
				y = (double)j;
				taux.data[0][j][i] = 1.*Math.cos( 2.*Math.PI*y/((double)taux.ny));
			}
		}

		for (int j=0; j < tauy.ny; j++) {
			for (int i=0; i < tauy.nx; i++) {      
				tauy.data[0][j][i] = 0.;
			}
		}
		
	}

	//-----------------------------------------------------------------------
	  
	public void calculate(HashMap<String, GriddedVariable> variables, 
			              HashMap<String, double[][][]> derivatives, 
			              Parameters params) {

		double[][][] du_dt = derivatives.get("u");
		double[][][] dv_dt = derivatives.get("v");
		
		GriddedVariable h = variables.get("h");
		GriddedVariable u = variables.get("u");
		GriddedVariable v = variables.get("v");
	    
		for (int j=0; j < u.ny; j++) {
			for (int i=0; i < u.nx; i++) {      
				du_dt[0][j][i] += taux.data[0][j][i]/params.rho_0/400.;
			}
		}
	        
		for (int j=0; j < v.ny; j++) {
			for (int i=0; i < v.nx; i++) {      
				dv_dt[0][j][i] += tauy.data[0][j][i]/params.rho_0/400.;
			}
		}
	    
	}
	  
}

//=================================================

class Pressure {
	  
	// Calculates hydrostatic pressure force

	Pressure(HashMap<String, GriddedVariable> variables, 
			 HashMap<String, double[][][]> derivatives, 
			 Parameters params) {
		
	}
	  
	//-------------------------------------------------------------------------------

	public void calculate(HashMap<String, GriddedVariable> variables, 
			              HashMap<String, double[][][]> derivatives, 
			              Parameters params) {
		double[][][] data_dp_dx, data_dp_dy;

		GriddedVariable eta = variables.get("eta");
		GriddedVariable press = variables.get("pressure");
	    GriddedVariable u = variables.get("u");
	    GriddedVariable v = variables.get("v");

	    double[][][] du_dt = derivatives.get("u");
	    double[][][] dv_dt = derivatives.get("v");
	    
	    for (int j=0; j < press.ny; j++) {
	    	for (int i=0; i < press.nx; i++) {      
	    		press.data[0][j][i] = params.rho[0]*params.g*eta.data[0][j][i];
	    	}
	    }

	    for (int k=1; k < press.nz; k++) {
	    	for (int j=0; j < press.ny; j++) {
	    		for (int i=0; i < press.nx; i++) {      
	    			press.data[k][j][i] = press.data[k-1][j][i] + (params.rho[k]-params.rho[k-1])*params.g*eta.data[k][j][i];
	    		}
	    	}
	    }
	        
	    data_dp_dx = press.d_dx().data;
	    for (int k=0; k < u.nz; k++) {
	    	for (int j=0; j < u.ny; j++) {
	    		for (int i=0; i < u.nx; i++) {      
	    			du_dt[k][j][i] += -data_dp_dx[k][j][i]/params.rho_0;
	    		}
	    	}
	    }

	    data_dp_dy = press.d_dy().data;
	    for (int k=0; k < v.nz; k++) {
	    	for (int j=0; j < v.ny; j++) {
	    		for (int i=0; i < v.nx; i++) {      
	    			dv_dt[k][j][i] += -data_dp_dy[k][j][i]/params.rho_0;
	    		}
	    	}
	    }
	}

	  
}

//===================================================================

class Viscosity {
	  
	// Calculates eddy viscosity force

	Viscosity(HashMap<String, GriddedVariable> variables, 
			  HashMap<String, double[][][]> derivatives, 
			  Parameters params) {
		
	}
	    
	public void calculate(HashMap<String, GriddedVariable> variables, 
			              HashMap<String, double[][][]> derivatives, 
			              Parameters params) {
		GriddedVariable u = variables.get("u");
		GriddedVariable v = variables.get("v");

		double[][][] du_dt = derivatives.get("u");
		double[][][] dv_dt = derivatives.get("v");

		double[][][] data_d2u_dx2, data_d2u_dy2;
		double[][][] data_d2v_dx2, data_d2v_dy2;
	                        
		data_d2u_dx2 = u.d2_dx2().data;
		data_d2u_dy2 = u.d2_dy2().data;
		for (int k=0; k < u.nz; k++) {
			for (int j=0; j < u.ny; j++) {
				for (int i=0; i < u.nx; i++) {      
					du_dt[k][j][i] += (data_d2u_dx2[k][j][i] + data_d2u_dy2[k][j][i])*params.kappa;      
				}
			}
		}

		data_d2v_dx2 = v.d2_dy2().data;
		data_d2v_dy2 = v.d2_dx2().data;
		for (int k=0; k < v.nz; k++) {
			for (int j=0; j < v.ny; j++) {
				for (int i=0; i < v.nx; i++) {      
					dv_dt[k][j][i] += (data_d2v_dx2[k][j][i] + data_d2v_dy2[k][j][i])*params.kappa;
				}
			}
		}
		
	}
	
}
